package com.lichao.sink;

import com.lichao.source.Event;
import com.lichao.utils.ReadProperties;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.util.*;

public class SinkToClickhouse extends RichSinkFunction<Event> {

    private static final int BATCH_SIZE = 2;
    public Map<String, Connection> connectionMap = new LinkedHashMap<>();
    private final List<Event> batchData = new ArrayList<>(BATCH_SIZE);

    private static final int TABLE_SIZE = 0;

    @Override
    public void open(Configuration parameters) throws Exception {
        System.out.println("open 函数调用");
        super.open(parameters);
        Class.forName("com.mysql.jdbc.Driver");
        connectionMap.put("mysql1", DriverManager.getConnection(ReadProperties.MYSql1, "root", "123456"));
        connectionMap.put("mysql2", DriverManager.getConnection(ReadProperties.MYSql2, "root", "123456"));
        connectionMap.put("mysql3", DriverManager.getConnection(ReadProperties.MYSql3, "root", "123456"));
    }

    @Override
    public void invoke(Event value, Context context) throws Exception {
        batchData.add(value);
        if (batchData.size() >= BATCH_SIZE) {
            flushBatch();
        }
    }

    @Override
    public void close() throws Exception {
        System.out.println("close 函数调用");
        super.close();
    }

    // Map<数据库, Map<表名, PreparedStatement>>
    public Map<String, Map<String, PreparedStatement>> statementMap = new LinkedHashMap<>();


    private static final int MAX_RETRIES = 3; // 最大重试次数

    private void flushBatch() throws Exception {
        System.out.println("invoke 函数调用");
        // 循环数据连接数据
        Set<String> tableSet = ReadProperties.tableSet;
        if (TABLE_SIZE != tableSet.size()) {
            statementMap.clear();
            for (String mysql : connectionMap.keySet()) {
                Connection connection = connectionMap.get(mysql);
                Map<String, PreparedStatement> table = new LinkedHashMap<>();
                for (String tableName : tableSet) {
                    String sql = "INSERT INTO " + tableName + " (url) VALUES (?)";
                    table.put(tableName, connection.prepareStatement(sql));
                }
                statementMap.put(mysql, table);
            }
        }

        int retryCount = 0;
        while (retryCount < MAX_RETRIES && !batchData.isEmpty()) {
            try {
                Map.Entry<String, Map<String, PreparedStatement>> next = getNext();
                Map<String, PreparedStatement> values = next.getValue();

                for (Event value : batchData) {
                    PreparedStatement statement = values.get(value.user);
                    if (statement == null) {
                        throw new IllegalStateException("No prepared statement found for user: " + value.user);
                    }
                    statement.setString(1, value.url);
                    statement.addBatch();
                }

                // 在这里，你可以遍历statementMap中的每个PreparedStatement并执行executeBatch()
                // 但为了简化示例，我们假设只有一个entry需要被处理
                for (Map<String, PreparedStatement> table : statementMap.values()) {
                    for (PreparedStatement stmt : table.values()) {
                        stmt.executeBatch();
                        stmt.clearBatch(); // 清除批处理命令
                    }
                }

                batchData.clear(); // 清除已处理的批处理数据
                retryCount = MAX_RETRIES; // 跳出循环

            } catch (Exception e) {
                retryCount++; // 重试次数加一
                if (retryCount >= MAX_RETRIES) {
                    throw e; // 达到最大重试次数，重新抛出异常
                }
                // 这里可以添加一些日志记录或处理逻辑
                System.err.println("Failed to execute batch, retrying... (" + retryCount + "/" + MAX_RETRIES + ")");
            }
        }


    }


    private int currentIndex = 0; // 当前索引，用于轮询

    /**
     * 获取下一个元素，如果到达末尾则从头开始
     *
     * @return 下一个键值对
     */
    public Map.Entry<String, Map<String, PreparedStatement>> getNext() {
        // 检查是否到达末尾
        if (currentIndex >= statementMap.size()) {
            currentIndex = 0; // 重置索引
        }

        // 获取当前索引的元素
        Map.Entry<String, Map<String, PreparedStatement>> entry = null;
        int i = 0;
        for (Map.Entry<String, Map<String, PreparedStatement>> e : statementMap.entrySet()) {
            if (i == currentIndex) {
                entry = e;
                break;
            }
            i++;
        }

        // 移动到下一个索引
        currentIndex++;
        return entry;
    }

}
